package com.ficus.road.maintain.util;


import lombok.extern.slf4j.Slf4j;

import java.util.Comparator;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.BiConsumer;

@Slf4j
public class TimeoutMap<K, V> {
    private TreeMap<K, DataEntry<V>> dataMap;
    private TreeSet<TimeoutEntry<K>> timeoutSet;
    private ReentrantLock lock;
    private Condition condition;
    private BiConsumer<K, V> timeoutConsumer;
    private int timeoutInMs;

    public TimeoutMap(int timeoutInMs,
                      BiConsumer<K, V> timeoutConsumer) {
        this.dataMap = new TreeMap<>();
        this.timeoutSet = new TreeSet<>(new DefaultTimeoutEntryComparator<>());
        this.lock = new ReentrantLock();
        this.condition = this.lock.newCondition();
        this.timeoutConsumer = timeoutConsumer;
        this.timeoutInMs = timeoutInMs;
        new Thread(this::TimeoutLoop).start();
    }

    public TimeoutMap(Comparator<K> comparator,
                      int timeoutInMs,
                      BiConsumer<K, V> timeoutConsumer) {
        this.dataMap = new TreeMap<>(comparator);
        this.timeoutSet = new TreeSet<>(new TimeoutEntryComparator<>(comparator));
        this.lock = new ReentrantLock();
        this.condition = this.lock.newCondition();
        this.timeoutConsumer = timeoutConsumer;
        this.timeoutInMs = timeoutInMs;
        new Thread(this::TimeoutLoop).start();
    }

    public int size() {
        int count = 0;
        lock.lock();
        try {
            count = dataMap.size();
        } finally {
            lock.unlock();
        }
        return count;
    }

    public boolean isEmpty() {
        boolean empty = false;
        lock.lock();
        try {
            empty = dataMap.isEmpty();
        } finally {
            lock.unlock();
        }
        return empty;
    }

    public V put(K key, V value) {
        TimeoutEntry<K> timeoutEntry = new TimeoutEntry<>(key);
        DataEntry<V> dataEntry = new DataEntry<>(value);
        DataEntry<V> result = null;
        boolean empty = false;

        lock.lock();
        try {
            empty = dataMap.isEmpty();
            result = dataMap.put(key, dataEntry);
            dataEntry.timeoutTimeInMs = System.currentTimeMillis() + timeoutInMs;
            if (result == null) {
                timeoutEntry.timeoutTimeInMs = dataEntry.timeoutTimeInMs;
                timeoutSet.add(timeoutEntry);
                if (empty) {
                    condition.signal();
                }
            } else {
                // 替换timeoutSet中的节点
                timeoutEntry.timeoutTimeInMs = result.timeoutTimeInMs;
                timeoutSet.remove(timeoutEntry);
                timeoutEntry.timeoutTimeInMs = dataEntry.timeoutTimeInMs;
                timeoutSet.add(timeoutEntry);
            }
        } finally {
            lock.unlock();
        }
        return ((result == null) ? null : result.value);
    }

    public V remove(K key) {
        DataEntry<V> result = null;
        TimeoutEntry<K> timeoutEntry = new TimeoutEntry<>(key);
        lock.lock();
        try {
            result = dataMap.remove(key);
            if (result != null) {
                timeoutEntry.timeoutTimeInMs = result.timeoutTimeInMs;
                timeoutSet.remove(timeoutEntry);
            }
        } finally {
            lock.unlock();
        }
        return ((result == null) ? null : result.value);
    }

    public V get(K key) {
        DataEntry<V> result = null;
        lock.lock();
        try {
            result = dataMap.get(key);
        } finally {
            lock.unlock();
        }
        return ((result == null) ? null : result.value);
    }

    public V getOrDefault(K key, V defaultValue) {
        TimeoutEntry<K> timeoutEntry = new TimeoutEntry<>(key);
        DataEntry<V> dataEntry = new DataEntry<>(defaultValue);
        DataEntry<V> result = null;

        lock.lock();
        try {
            result = dataMap.get(key);
            if (result == null) {
                result = dataEntry;
                boolean empty = dataMap.isEmpty();
                dataMap.put(key, dataEntry);
                dataEntry.timeoutTimeInMs = System.currentTimeMillis() + timeoutInMs;
                timeoutEntry.timeoutTimeInMs = dataEntry.timeoutTimeInMs;
                timeoutSet.add(timeoutEntry);
                if (empty) {
                    condition.signal();
                }
            }
        } finally {
            lock.unlock();
        }
        return result.value;
    }

    private void TimeoutLoop() {
        K timeoutKey = null;
        V timeoutValue = null;
        lock.lock();
        while (true) {
            try {
                while (dataMap.isEmpty()) {
                    condition.await();
                }
            } catch (Exception e) {
                log.error("", e);
            }
            long now = System.currentTimeMillis();
            TimeoutEntry<K> timeoutEntry = timeoutSet.first();
            if (now >= timeoutEntry.timeoutTimeInMs) {
                timeoutKey = timeoutEntry.key;
                timeoutSet.remove(timeoutEntry);
                timeoutValue = dataMap.remove(timeoutKey).value;
                lock.unlock();
                try {
                    timeoutConsumer.accept(timeoutKey, timeoutValue);
                } catch (Exception e) {
                    log.error("", e);
                }
            } else {
                long delayInMs = timeoutEntry.timeoutTimeInMs - now;
                lock.unlock();
                try {
                    Thread.sleep(delayInMs);
                } catch (Exception e) {
                    log.error("", e);
                }
            }
            lock.lock();
        }
    }

    static private class DataEntry<E> {
        E value;
        long timeoutTimeInMs;
        DataEntry(E inputValue) {
            this.value = inputValue;
        }
    }

    static private class TimeoutEntry<E>  {
        E key;
        long timeoutTimeInMs;
        TimeoutEntry(E inputKey) {
            this.key = inputKey;
        }
    }

    static private class TimeoutEntryComparator<E> implements Comparator<TimeoutEntry<E>>{
        private Comparator<E> comparator;

        TimeoutEntryComparator(Comparator<E> comparator) {
            this.comparator = comparator;
        }

        @Override
        public int compare(TimeoutEntry<E> entry1, TimeoutEntry<E> entry2) {
            int c = Long.compare(entry1.timeoutTimeInMs, entry2.timeoutTimeInMs);
            if (c == 0) {
                return comparator.compare(entry1.key, entry2.key);
            } else {
                return c;
            }
        }
    }

    static private class DefaultTimeoutEntryComparator<E> implements Comparator<TimeoutEntry<E>>{

        @Override
        public int compare(TimeoutEntry<E> entry1, TimeoutEntry<E> entry2) {
            int c = Long.compare(entry1.timeoutTimeInMs, entry2.timeoutTimeInMs);
            if (c == 0) {
                return ((Comparable<E>)(entry1.key)).compareTo(entry2.key);
            } else {
                return c;
            }
        }
    }
}
